# Autor chenfeng
#!/usr/bin/env Python
# coding=utf-8

from paddlenlp.datasets import load_dataset
import json


# 读取标签文件并创建label_map
def get_label_map(filename):
    with open(filename, "r", encoding="utf8") as f:
        label_list = json.load(f)

    id2label = dict([(idx, label) for idx, label in enumerate(label_list)])
    label2id = dict([(label, idx) for idx, label in enumerate(label_list)])
    return id2label, label2id


id2slot, slot2id = get_label_map('slot_labels.json')
id2intent, intent2id = get_label_map('intent_labels.json')

intent_weight = [1] * len(intent2id)


# 根据本地文件格式定义数据读取生成器
def read(filename):
    with open(filename, "r", encoding="utf8") as f:
        data = json.load(f)
    for entry in data:
        yield {
            'words': entry[0],
            'slots': entry[1],
            'intents': entry[2],
            'history': entry[4],
        }


# 将生成器传入load_dataset
train_ds = load_dataset(read, filename='train.json', lazy=False)
dev_ds = load_dataset(read, filename='test.json', lazy=False)

for idx in range(2):
    print(train_ds[idx]['words'])
    print(train_ds[idx]['slots'])
    print(train_ds[idx]['intents'])
    print(train_ds[idx]['history'])
    print()
